import torch
import cv2 as cv
import numpy as np


def decode_box(x, device=torch.device('cpu')):
    """输出网络的实际预测结果
    预测结果为:横坐标,纵坐标,宽度,高度,置信度
    :return 相对每个网格的横纵坐标和宽高,再乘32为实际值
    """
    x = [x[..., :300], x[..., 300:1500], x[..., 1500:]]
    x = [x[0].reshape(-1, 5, 15, 20), x[1].reshape(-1, 5, 30, 40), x[2].reshape(-1, 5, 60, 80)]
    predictions = []
    for i in range(3):
        predict = x[i].permute(0, 2, 3, 1)  # 批大小,高,宽,预测结果 N,H,W,P
        img_height, img_width = predict.shape[1], predict.shape[2]
        grid_x = torch.arange(0, img_width).repeat(img_height, 1)
        grid_y = torch.arange(0, img_height).repeat(img_width, 1).transpose(1, 0)
        grid_x = grid_x.to(device)
        grid_y = grid_y.to(device)
        predict[..., 0] = torch.sigmoid(predict[..., 0]) + grid_x  # 横坐标
        predict[..., 1] = torch.sigmoid(predict[..., 1]) + grid_y  # 纵坐标
        predict[..., 2] = 7 * torch.exp(predict[..., 2])  # 宽度
        predict[..., 3] = 10 * torch.exp(predict[..., 3])  # 高度
        predict[..., 4] = torch.sigmoid(predict[..., 4])  # 置信度
        predictions.append(predict)
    return tuple(predictions)


def image_generate_conf(img, label):
    """在图像上显示置信度
    :param img:网络输入的原图像
    :param label:网络经解码后的输出
    """
    conf = '{:.4f}'.format(label)
    img = cv.putText(img, conf, (160, 160), cv.FONT_HERSHEY_PLAIN, 1.2, (0, 0, 255), 1)
    return img


def image_generate_loc(img, label):
    """在图像上显示结果
    :param img:网络输入的原图像
    :param label:网络经解码后的输出
    """
    stride = 32, 16, 8
    cnt = [0, 0, 0]
    pred = []
    for fm in range(3):
        pred.append(label[fm].view(-1, 5).numpy())
        cnt[fm] = pred[fm].shape[0]
    pred = np.vstack(pred)
    index = pred[:, 4].argmax()
    pred = pred[index]
    if pred[4] < 0.4:
        return img
    corner = pred[:4].copy()
    corner[0] = pred[0] - pred[2]/2
    corner[1] = pred[1] - pred[3]/2
    corner[2] = pred[0] + pred[2]/2
    corner[3] = pred[1] + pred[3]/2
    fm = 0
    while index > 0:
        index -= cnt[fm]
        fm += 1
    fm -= 1
    corner = np.round(corner*stride[fm]).astype('int32')
    img = cv.rectangle(img, (corner[0], corner[1]), (corner[2], corner[3]), color=(0, 0, 255), thickness=2)
    pred_conf = '{:.4f}'.format(pred[4])
    img = cv.putText(img, pred_conf, (corner[0], corner[1]), cv.FONT_HERSHEY_PLAIN, 1.2, (0, 0, 255), 1)
    return img
